#!/usr/bin/env python3

# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
import numpy as np
from scene import Scene
import os
from tqdm import tqdm
from os import makedirs
from gaussian_renderer import render
import torchvision
from utils.general_utils import safe_state
from argparse import ArgumentParser
from arguments import ModelParams, PipelineParams, get_combined_args
from gaussian_renderer import GaussianModel
import copy, time
import imageio


to8b = lambda x: x.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()

def interpolate_views(view0, view1, length):
    #view0 source view
    #view1 target view 
    results = []
    for i in range(length):
        weight = (i+1) / length
        view = copy.deepcopy(view0)
        view.FoVy = view0.FoVy * (1-weight) + view1.FoVy * weight
        view.FoVx = view0.FoVx * (1-weight) + view1.FoVx * weight
        view.camera_center = view0.camera_center * (1-weight) + view1.camera_center * weight
        view.projection_matrix      = view0.projection_matrix * (1-weight) + view1.projection_matrix * weight
        view.world_view_transform   = view0.world_view_transform * (1-weight) + view1.world_view_transform * weight
        view.full_proj_transform    = (view.world_view_transform.unsqueeze(0).bmm(view.projection_matrix.unsqueeze(0))).squeeze(0)
        results.append(view)
        del view
    return results

def generate_multi_views(views, length=60):
    generated_views = []
    for i in tqdm(range(len(views)-1), desc='Interpolate views ...'):
        views_temp = interpolate_views(views[i], views[i+1], length)
        generated_views.extend(views_temp)
    return generated_views    

         
def render_video(dataset: ModelParams, iteration: int, pipeline: PipelineParams):
    with torch.no_grad():
        gaussians = GaussianModel(dataset.sh_degree)
        scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)
        bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
        model_path = dataset.model_path
        
        root, name = model_path.split('/')
        train_views = scene.getTrainCameras()
        if name.startswith('bicycle'):
            base_view_idxs = [0, 6, 12, 18, 24]
        elif name.startswith('bonsai'):      
            base_view_idxs = [21, 46, 71, 94, 120]
        elif name.startswith('counter'):      
            base_view_idxs = [21, 46, 71, 94, 120]
        elif name.startswith('flowers'):      
            base_view_idxs = [0, 4, 8, 12, 20]
        elif name.startswith('garden'):      
            base_view_idxs = [0, 4, 10, 16, 20]
        elif name.startswith('kitchen'):      
            base_view_idxs = [0, 6, 16, 16, 66]
        elif name.startswith('room'):      
            base_view_idxs = [0, 4, 19]
        elif name.startswith('stump'):      
            base_view_idxs = [0, 2, 4, 6, 8]
        elif name.startswith('treehill'):      
            base_view_idxs = [64, 66, 70, 74, 80]
            
        else:
            raise TypeError(f'No type nameed {name}!')
            
              
        base_views = [train_views[i] for i in range(len(train_views)) if i in base_view_idxs]
        length_view = 120
        generated_views = generate_multi_views(base_views, length_view)
        
        
        render_video_out = imageio.get_writer(f'{root}/{name}.mp4', mode='I', fps=30,codec='libx264',quality=10.0)#
        
        t_list = []
        save_dir = os.path.join(model_path, 'video_images')
        os.makedirs(save_dir, exist_ok=True)
        for idx, view in enumerate(tqdm(generated_views, desc="Rendering progress")):
            torch.cuda.synchronize(); t0 = time.time()        
            image = render(view, gaussians, pipeline, background)["render"]
            torch.cuda.synchronize(); t1 = time.time()
            np_image = to8b(image)
            render_video_out.append_data(np_image)
            # torchvision.utils.save_image(image, os.path.join(save_dir, f"video_image_{idx:05d}.png"))
            t_list.append(t1-t0)
            del image, np_image
        t = np.array(t_list[5:])
        fps = 1.0 / t.mean()
        print(f'Test FPS: \033[1;35m{fps:.5f}\033[0m')
        render_video_out.close()


if __name__ == "__main__":
    # Set up command line argument parser
    parser = ArgumentParser(description="Testing script parameters")
    model = ModelParams(parser, sentinel=True)
    pipeline = PipelineParams(parser)
    parser.add_argument("--iteration", default=-1, type=int)
    parser.add_argument("--skip_train", action="store_true")
    parser.add_argument("--skip_test", action="store_true")
    parser.add_argument("--quiet", action="store_true")
    args = get_combined_args(parser)
    print("Rendering " + args.model_path)

    # Initialize system state (RNG)
    safe_state(args.quiet)

    render_video(model.extract(args), args.iteration, pipeline.extract(args))